package com.ruoyi.web.service.impl;


import com.plexpt.chatgpt.entity.chat.Message;
import com.ruoyi.common.ai.dto.AIChatOpenSseDTO;
import com.ruoyi.common.base.chat.domain.Chat;
import com.ruoyi.common.base.chat.domain.ChatMessage;
import com.ruoyi.common.base.chat.dto.ChatMessageAddDTO;
import com.ruoyi.common.base.chat.dto.ChatMessageAppListCondition;
import com.ruoyi.common.base.chat.mapper.ChatMapper;
import com.ruoyi.common.base.chat.mapper.ChatMessageMapper;
import com.ruoyi.common.base.chat.mapstruct.ChatMessageStruct;
import com.ruoyi.common.base.chat.vo.ChatMessageAddRes;
import com.ruoyi.common.constant.SysConfigConstants;
import com.ruoyi.common.core.page.TableDataInfo;
import com.ruoyi.common.customize.util.CommonUtil;
import com.ruoyi.common.customize.util.chat.AIChatUtils;
import com.ruoyi.common.utils.DateUtils;
import com.ruoyi.system.service.ISysConfigService;
import com.ruoyi.web.service.IChatMessageService;
import lombok.AllArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import java.util.*;

/**
 * 会话记录Service业务层处理
 *
 * @author gdk
 * @date 2023-05-15
 */
@Service
@AllArgsConstructor
public class ChatMessageServiceImpl implements IChatMessageService {
    private final ChatMessageMapper chatMessageMapper;
    private final ChatMapper chatMapper;

    private final CommonUtil commonUtil;
    private final ISysConfigService sysConfigService;

    /**
     * 查询会话记录
     *
     * @param id 会话记录主键
     * @return 会话记录
     */
    @Override
    public ChatMessage selectChatMessageById(Long id) {
        return chatMessageMapper.selectChatMessageById(id);
    }

    /**
     * 查询会话记录列表
     *
     * @param chatMessage 会话记录
     * @return 会话记录
     */
    @Override
    public List<ChatMessage> selectChatMessageList(ChatMessage chatMessage) {
        return chatMessageMapper.selectChatMessageList(chatMessage);
    }

    /**
     * 新增会话记录
     *
     * @param chatMessage 会话记录
     * @return 结果
     */
    @Override
    public int insertChatMessage(ChatMessage chatMessage) {
        chatMessage.setCreateTime(DateUtils.getNowDate());
        return chatMessageMapper.insertChatMessage(chatMessage);
    }

    /**
     * 修改会话记录
     *
     * @param chatMessage 会话记录
     * @return 结果
     */
    @Override
    public int updateChatMessage(ChatMessage chatMessage) {
        chatMessage.setUpdateTime(DateUtils.getNowDate());
        return chatMessageMapper.updateChatMessage(chatMessage);
    }

    /**
     * 批量删除会话记录
     *
     * @param ids 需要删除的会话记录主键
     * @return 结果
     */
    @Override
    public int deleteChatMessageByIds(Long[] ids) {
        return chatMessageMapper.deleteChatMessageByIds(ids);
    }

    /**
     * 删除会话记录信息
     *
     * @param id 会话记录主键
     * @return 结果
     */
    @Override
    public int deleteChatMessageById(Long id) {
        return chatMessageMapper.deleteChatMessageById(id);
    }

    /**
     * 记录消息
     *
     * @param dto dto
     * @return {@link ChatMessageAddRes}
     * @author gdk
     * @since 2023/5/16 20:18
     */
    @Override
    @Transactional(rollbackFor = Throwable.class)
    public ChatMessageAddRes add(ChatMessageAddDTO dto) {
        // 获取或创建会话
        Chat chat = getOrAddChat(dto.getChatId(), dto.getPrompt());
        boolean isNewChat = !chat.getId().equals(dto.getChatId());

        ChatMessage message = ChatMessageStruct.INSTANCE.addDTO2Entity(dto);
        commonUtil.initCreateAbout(message);
        // 不是新会话，有id
        if (!isNewChat && !Objects.isNull(dto.getId())) {
            // 可能是[重新回答]，尝试修改回答
            if (chatMessageMapper.retry(message) == 1) {
                ChatMessageAddRes res = new ChatMessageAddRes();
                res.setChatId(chat.getId());
                res.setMsgId(message.getId());
                return res;

            } else {
                // 会话记录不存在/已被删除
                dto.setId(commonUtil.getId());
            }
        }

        // 创建会话记录
        message.setChatId(chat.getId());
        message.setRetryCount(0);
        if (chatMessageMapper.insertChatMessage(message) != 1) {
            CommonUtil.throwError();
        }

        // 组装返回结果
        ChatMessageAddRes res = new ChatMessageAddRes();
        res.setChatId(chat.getId());
        res.setMsgId(message.getId());
        if (isNewChat) {
            res.setTitle(chat.getTitle());
        }
        return res;
    }

    /**
     * 获取或创建会话
     *
     * @param chatId  会话id
     * @param content 用户消息
     * @return {@link Chat}
     * @author gdk
     * @since 2023/5/16 20:38
     */
    public Chat getOrAddChat(Long chatId, String content) {
        // chatId不为空，查询会话
        if (!Objects.isNull(chatId)) {
            Chat chat = new Chat();
            chat.setId(chatId);
            chat.setCreateBy(CommonUtil.getStrUserIdNotNull());
            chat = chatMapper.getOneLock(chat);
            // 已有会话，修改更新时间
            if (!Objects.isNull(chat)) {
                chat.setUpdateTime(new Date());
                if (chatMapper.updateChat(chat) != 1) {
                    CommonUtil.throwError();
                }
                return chat;
            }
        }

        // 会话数据尚未创建/丢失，创建新的会话
        Chat chat = new Chat();
        chat.setTitle(AIChatUtils.getChatTitle(content));
        commonUtil.initCreateAbout(chat);
        if (chatMapper.insertChat(chat) != 1) {
            CommonUtil.throwError();
        }

        return chat;
    }

    /**
     * 消息记录
     *
     * @param condition 条件
     * @return {@link TableDataInfo}
     * @author gdk
     * @since 2023/5/16 17:08
     */
    @Override
    public TableDataInfo appList(ChatMessageAppListCondition condition) {
        condition.setCreateBy(CommonUtil.getUserIdNotNull());
        List<ChatMessage> list = chatMessageMapper.appList(condition);

        return new TableDataInfo(list, ChatMessageStruct.INSTANCE.list2AppVOList(list));
    }

    /**
     * 初始化发送给AI的消息列表
     *
     * @param dto dto
     * @return {@link List< Message>}
     * @author gdk
     * @since 2023/5/22 14:55
     */
    @Override
    public List<Message> initMessageList(AIChatOpenSseDTO dto) {
        // 确保是本人数据
        Long userId = CommonUtil.getUserIdNotNull();
        int pageSize = Integer.parseInt(sysConfigService.selectConfigByKey(SysConfigConstants.AI_CHAT_MAX_CONTEXT, "10"));
        int maxContextLength = Integer.parseInt(sysConfigService.selectConfigByKey(SysConfigConstants.AI_CHAT_MAX_CONTEXT_LENGTH, "1000"));

        // chatId为空，没有上下文
        if (Objects.isNull(dto.getChatId())) {
            return Arrays.asList(Message.of(dto.getPrompt()));
        }

        // 根据chatId、messageId查询上下文，messageId存在表明是重试，需要将该消息剔除
        List<ChatMessage> chatMessageList = chatMessageMapper.listAIChatContext(userId,
                dto.getChatId(),
                dto.getChatMessageId(),
                pageSize);
        // 没有上下文
        if (CollectionUtils.isEmpty(chatMessageList)) {
            return Arrays.asList(Message.of(dto.getPrompt()));
        }

        List<Message> messageList = new ArrayList<>(chatMessageList.size() + 1);
        // 用户此次提示词
        messageList.add(Message.of(dto.getPrompt()));
        // 遍历
        int totalLength = 0;
        for (ChatMessage item : chatMessageList) {
            // AI回复
            Message assistantMessage = getMessage(Message.Role.ASSISTANT, item.getReply(), totalLength, maxContextLength);
            messageList.add(assistantMessage);
            if ((totalLength += assistantMessage.getContent().length()) == maxContextLength) {
                break;
            }

            // 用户提示词
            Message userMessage = getMessage(Message.Role.USER, item.getPrompt(), totalLength, maxContextLength);
            messageList.add(userMessage);
            if ((totalLength += userMessage.getContent().length()) == maxContextLength) {
                break;
            }
        }

        // 翻转
        Collections.reverse(messageList);
        return messageList;
    }

    /**
     * 处理上下文内容
     *
     * @param role        角色
     * @param content     内容
     * @param totalLength 总长度
     * @param maxLength   最大长度
     * @return {@link Message}
     * @author gdk
     * @since 2023/5/22 15:47
     */
    private Message getMessage(Message.Role role, String content, int totalLength, int maxLength) {
        totalLength += content.length();
        if (totalLength > maxLength) {
            int endIndex = content.length() - (totalLength - maxLength);
            content = content.substring(0, endIndex);
        }

        return Message.builder()
                .role(role.getValue())
                .content(content)
                .build();
    }
}
